import hmac
import logging
import os

from flask import Flask, request, jsonify, Response
from model import YOLOv8Model


logger = logging.getLogger(__name__)

_server = Flask(__name__)
MODEL_CLASS = YOLOv8Model


@_server.route('/predict', methods=['POST'])
def _predict():
    """
    Predict tasks

    Example request:
    request = {
            'tasks': tasks,
            'model_version': model_version,
            'project': '{project.id}.{int(project.created_at.timestamp())}',
            'label_config': project.label_config,
            'params': {
                'login': project.task_data_login,
                'password': project.task_data_password,
                'context': context,
            },
        }

    @return:
    Predictions in LS format
    """
    data = request.json
    logger.debug(data)
    tasks = data.get('tasks')
    label_config = data.get('label_config')
    project = str(data.get('project'))
    project_id = project.split('.', 1)[0] if project else None
    params = data.get('params', {})
    context = params.pop('context', {})

    predictions = []
    for task in tasks:
        model = MODEL_CLASS(project_id=project_id,label_config=label_config)
        regions = model.predict(task, context=context, **params)
        all_scores = [region["score"] for region in regions if "score" in region]
        avg_score = sum(all_scores) / max(len(all_scores), 1)
        prediction = {
            "result": regions,
            "score": avg_score,
            "model_version": 'yolo11n_dnf',
        }
        predictions.append(prediction)

    # model.use_label_config(label_config)


    return jsonify({'results': predictions})


@_server.route('/setup', methods=['POST'])
def _setup():
    data = request.json
    project_id = data.get('project').split('.', 1)[0]
    label_config = data.get('schema')
    extra_params = data.get('extra_params')
    return jsonify({'model_version': 'yolo11n_dnf'})


TRAIN_EVENTS = (
    'ANNOTATION_CREATED',
    'ANNOTATION_UPDATED',
    'ANNOTATION_DELETED',
    'START_TRAINING'
)


@_server.route('/webhook', methods=['POST'])
def webhook():
    data = request.json
    event = data.pop('action')
    if event not in TRAIN_EVENTS:
        return jsonify({'status': 'Unknown event'}), 200
    project_id = str(data['project']['id'])
    label_config = data['project']['label_config']
    model = MODEL_CLASS(project_id, label_config=label_config)
    result = model.fit(event, data)

    try:
        response = jsonify({'result': result, 'status': 'ok'})
    except Exception as e:
        response = jsonify({'error': str(e), 'status': 'error'})

    return response, 201


@_server.route('/health', methods=['GET'])
@_server.route('/', methods=['GET'])
def health():
    return jsonify({'status': 'UP'})


@_server.route('/metrics', methods=['GET'])
def metrics():
    return jsonify({})


@_server.errorhandler(FileNotFoundError)
def file_not_found_error_handler(error):
    logger.warning('Got error: ' + str(error))
    return str(error), 404


@_server.errorhandler(AssertionError)
def assertion_error(error):
    logger.error(str(error), exc_info=True)
    return str(error), 500


@_server.errorhandler(IndexError)
def index_error(error):
    logger.error(str(error), exc_info=True)
    return str(error), 500


def safe_str_cmp(a, b):
    return hmac.compare_digest(a, b)




@_server.before_request
def log_request_info():
    logger.debug('Request headers: %s', request.headers)
    logger.debug('Request body: %s', request.get_data())


@_server.after_request
def log_response_info(response):
    logger.debug('Response status: %s', response.status)
    logger.debug('Response headers: %s', response.headers)
    logger.debug('Response body: %s', response.get_data())
    return response

if __name__ == '__main__':
    _server.run(host='0.0.0.0', port=9090)